import numpy as np
from ortools.linear_solver import pywraplp
from collections import defaultdict
import math
import random
import time
from typing import List, Dict, Tuple


class WeaponSystem:
    """武器系统基类"""

    def __init__(self, name: str, weapon_type: str, position: Tuple[float, float]):
        self.name = name
        self.weapon_type = weapon_type
        self.position = position  # (经度, 纬度)
        self.damage_status = 1.0  # 损伤状态 (0.0-1.0)

    def set_radar_params(self, radar_range: float, frequency: float, power: float):
        """设置雷达参数"""
        self.radar_params = {
            'range': radar_range,
            'frequency': frequency,
            'power': power
        }

    def set_aaw_systems(self, systems: List[Dict]):
        """设置防空系统参数"""
        self.aaw_systems = systems


class RedMissile(WeaponSystem):
    """红方导弹系统"""

    def __init__(self, name: str, missile_type: str, position: Tuple[float, float],
                 rcs: float, speed: float, altitude: float, warhead_yield: float,
                 quantity: int):
        """
        红方导弹初始化
        :param rcs: 雷达反射面积(m²)
        :param speed: 飞行速度(Mach)
        :param altitude: 飞行雷达反射面积高度(m)
        :param warhead_yield: 战斗部当量(kg TNT)
        :param quantity: 可用数量
        """
        super().__init__(name, f"missile_{missile_type}", position)
        self.rcs = rcs
        self.speed = speed  # 马赫数
        self.altitude = altitude
        self.warhead_yield = warhead_yield
        self.quantity = quantity
        self.cost = self._calculate_cost()

    def _calculate_cost(self):
        """计算导弹成本(简化模型)"""
        return (self.speed * 0.3 + self.warhead_yield * 0.1) * 1000000

    def get_effectiveness(self, target: 'BlueShip') -> float:
        """
        计算对特定目标的攻击效能
        :return: 0.0-1.0 效能值
        """
        # 基础命中概率
        base_hit_prob = 0.8 - (self.rcs * 0.05)  # 隐身性影响

        # 速度优势
        speed_bonus = min(self.speed / 5.0, 0.2)  # 最高奖励20%

        # 高度优势
        if self.altitude < 50:  # 超低空
            altitude_bonus = 0.15
        elif self.altitude > 10000:  # 高空
            altitude_bonus = 0.10
        else:
            altitude_bonus = 0

        return base_hit_prob + speed_bonus + altitude_bonus


class BlueShip(WeaponSystem):
    """蓝方舰船系统"""

    def __init__(self, name: str, ship_type: str, position: Tuple[float, float]):
        super().__init__(name, f"ship_{ship_type}", position)
        self.threat_level = 0.0
        self.defense_status = 1.0  # 防御系统状态

    def calculate_threat_level(self, missile_params: Dict) -> float:
        """
        计算目标威胁级别
        :param missile_params: 攻击导弹的参数
        :return: 威胁值 0.0-1.0
        """
        # 雷达特征计算
        detect_range = self._calculate_detect_range(missile_params['rcs'])
        radar_horizon = 25.0  # 雷达水平线距离(km)
        effective_range = min(detect_range, radar_horizon)
        radar_sig = effective_range / 400.0  # 归一化

        # 防空能力计算
        aaw_cap = self.calculate_aaw_capability()

        # 雷达探测能力权重40%，防空能力权重60%
        threat = radar_sig * 0.4 + aaw_cap * 0.6

        # 舰船类型权重
        ship_weights = {
            'carrier': 1.0, 'cruiser': 0.75, 'destroyer': 0.65,         #航空母舰、巡洋舰、驱逐舰、护卫舰、潜艇、补给舰
            'frigate': 0.45, 'submarine': 0.55, 'supply': 0.25
        }
        base_weight = ship_weights.get(self.weapon_type.replace('ship_', ''), 0.5)

        self.threat_level = threat * base_weight
        return self.threat_level

    def _calculate_detect_range(self, rcs: float) -> float:
        """计算雷达探测距离"""
        if not hasattr(self, 'radar_params'):
            return 0.0
        # 简化雷达方程
        return (self.radar_params['power'] *
                self.radar_params['range'] *
                math.sqrt(rcs) *
                (self.radar_params['frequency'] / 10))

    def calculate_aaw_capability(self) -> float:
        """计算防空能力"""
        if not hasattr(self, 'aaw_systems'):
            return 0.0

        aaw_weights = {'CIWS': 0.3, 'SAM': 0.6, 'ECM': 0.1}
        total_aaw_score = 0.0

        for system in self.aaw_systems:
            base_score = (aaw_weights[system['type']] *
                          system['count'])
                          #*math.sqrt(system['range'] / 100.0))
            #reload_penalty = math.exp(-system['reload_time'] / 30.0)
            total_aaw_score += base_score #* reload_penalty

        return min(total_aaw_score / 5.0, 1.0)

    def intercept_probability(self, missile: RedMissile) -> float:
        """
        计算拦截概率
        :param missile: 来袭导弹
        :return: 成功拦截概率
        """
        if not hasattr(self, 'aaw_systems'):
            return 0.0

        # 导弹速度影响
        speed_factor = 1.0 - min(0.5, missile.speed * 0.1)#min(1.0, 1 / (missile.speed * 0.25))

        # 导弹隐身影响
        stealth_factor =min(1.0, missile.rcs * 5) #min(1.0, missile.rcs * 10)

        # 防空能力基准
        base_prob = self.calculate_aaw_capability() * 0.7

        # 综合拦截概率
        prob = base_prob * speed_factor * stealth_factor * self.defense_status

        return min(prob, 0.95)  # 最大95%成功率


class FireAllocationSystem:
    """火力分配核心系统"""

    def __init__(self, config=None):
        self.red_missiles = []  # 红方导弹系统列表
        self.blue_ships = []  # 蓝方舰船列表
        self.cost_weight = 0.0    #变量
        self.min_missile_ratio =1
        self.diversity_weight = 1

    def set_allocation_strategy(self, diversity=0.4, min_types_ratio=0.4):
        """
        设置导弹组合策略
        :param diversity: 多样性权重 (0-1)
        :param min_types_ratio: 最少使用导弹类型比例 (0-1)
        """
        if 0 <= diversity <= 1:
            self.diversity_weight = diversity

        if 0 <= min_types_ratio <= 1:
            self.min_missile_ratio = min_types_ratio
    def add_red_missile(self, missile: RedMissile):
        """添加红方导弹系统"""
        self.red_missiles.append(missile)

    def add_blue_ship(self, ship: BlueShip):
        """添加蓝方舰船"""
        self.blue_ships.append(ship)

    def target_prioritization(self) -> List[BlueShip]:
        """目标优先级排序"""
        if not self.red_missiles:
            return sorted(self.blue_ships, key=lambda s: s.threat_level, reverse=True)

        # 获取典型导弹参数
        missile_params = {
            'rcs': self.red_missiles[0].rcs,
            'speed': self.red_missiles[0].speed,
            'altitude': self.red_missiles[0].altitude
        }

        # 计算每个目标威胁级别
        for ship in self.blue_ships:
            ship.calculate_threat_level(missile_params)

        return sorted(self.blue_ships, key=lambda s: s.threat_level, reverse=True)

    def allocate_red_firepower(self) -> Dict[str, Dict[str, int]]:
        """
        红方火力分配核心算法
        :return: 分配方案 {目标名称: {导弹类型: 数量}}
        """
        # 1. 目标优先级排序
        prioritized_targets = self.target_prioritization()
        # 2. 构建目标价值向量
        target_values = [t.threat_level for t in prioritized_targets]

        # 3. 获取导弹库存
        missile_inventory = {}
        for missile in self.red_missiles:
            missile_inventory[missile.name] = missile.quantity
        # 4. 优化分配 (使用整数线性规划)
        allocation_plan = self._optimize_allocation(prioritized_targets, target_values, missile_inventory)

        return allocation_plan

    def _optimize_allocation(self, targets, target_values, missile_inventory):
        # 创建求解器
        solver = pywraplp.Solver('DiverseAllocation', pywraplp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)

        # 导弹类型列表
        missile_types = list(missile_inventory.keys())

        # ===== 计算防空能力向量 =====
        aaw_capabilities = []
        for target in targets:
            # 计算每个目标的简化防空能力
            aaw_cap = target.calculate_aaw_capability()
            aaw_capabilities.append(aaw_cap)

        # ===== 决策变量 =====
        x = {}
        for i, target in enumerate(targets):
            for j, mt in enumerate(missile_types):
                x[i, j] = solver.IntVar(0, missile_inventory[mt],
                                        f'x_{target.name}_{mt}')
        y = {}
        for j, mt in enumerate(missile_types):
            y[j] = solver.BoolVar(f'use_{mt}')


        # ===== 约束条件 =====
        # 1. 导弹库存约束
        for j, mt in enumerate(missile_types):
            solver.Add(solver.Sum(x[i, j] for i in range(len(targets))) <= missile_inventory[mt])
        # 2. 导弹使用标记约束
        for j in range(len(missile_types)):
            solver.Add(y[j] * 10000 >= solver.Sum(x[i, j] for i in range(len(targets))))
            solver.Add(y[j] <= solver.Sum(x[i, j] for i in range(len(targets))))

        # 3. 最小导弹类型约束
        min_types = max(2, int(len(missile_types) * self.min_missile_ratio))
        solver.Add(solver.Sum(y[j] for j in range(len(missile_types))) >= min_types)
        # 4. 基于防空能力的导弹分配约束
        for i, target in enumerate(targets):
            total_missiles = solver.Sum(x[i, j] for j in range(len(missile_types)))
            aaw_cap = aaw_capabilities[i]
            # 根据防空能力确定最小和最大导弹数量
            min_missiles = int(1 + aaw_cap * 4)  # 防空能力越强，需要导弹越多
            max_missiles = int(3 + aaw_cap * 8)

            solver.Add(total_missiles >= min_missiles)
            solver.Add(total_missiles <= max_missiles)

            # 防空能力强的目标需要更多样化的导弹
            if aaw_cap > 0.6:  # 强防空目标
                solver.Add(solver.Sum(y[j] for j in range(len(missile_types))) >= min(3, len(missile_types)))

        # 5. 导弹类型均衡约束
        for j in range(len(missile_types)):
            min_alloc = int(missile_inventory[missile_types[j]] * 0.15)  # 每种导弹至少使用15%
            solver.Add(solver.Sum(x[i, j] for i in range(len(targets))) >= min_alloc)

        # ===== 目标函数 =====
        objective = solver.Objective()

        # 主目标: 毁伤效能净值
        for i, target in enumerate(targets):
            for j, mt in enumerate(missile_types):
                missile = next(m for m in self.red_missiles if m.name == mt)

                # 计算预期毁伤
                effectiveness = missile.get_effectiveness(target)
                expected_damage = effectiveness * missile.warhead_yield * target_values[i]

                # 成本调整
                cost_factor = self.cost_weight * missile.cost / 1000000

                # 防空能力调整 (防空能力越强，价值越低)
                aaw_factor = 1.0 - aaw_capabilities[i] * 0.3

                net_value = expected_damage * aaw_factor - cost_factor
                objective.SetCoefficient(x[i, j], net_value)

        # 多样性奖励
        diversity_bonus = self.diversity_weight * 2000
        for j in range(len(missile_types)):
            objective.SetCoefficient(y[j], diversity_bonus)

        objective.SetMaximization()

        # ===== 求解 =====
        solver.SetTimeLimit(3000)  # 3秒时间限制
        status = solver.Solve()

        # 处理结果
        allocation = defaultdict(lambda: defaultdict(int))
        if status == pywraplp.Solver.OPTIMAL or status == pywraplp.Solver.FEASIBLE:
            for i, target in enumerate(targets):
                for j, missile_name in enumerate(missile_inventory.keys()):
                    count = x[i, j].solution_value()
                    if count > 0:
                        allocation[target.name][missile_name] += int(round(count))
        else:
            print("优化器未找到最优解，使用后备方案")
            allocation = self._fallback_allocation(targets, target_values, missile_inventory)

        return dict(allocation)
    # 在FireAllocationSystem类中添加
    def set_cost_weight(self, weight: float):
        """
        设置成本权重系数
        :param weight: 0(不计成本)-1(极端保守)
        """
        if 0 <= weight <= 1:
            self.cost_weight = weight
        else:
            print(f"警告: 权重值{weight}超出范围，使用默认值0.6")
            self.cost_weight = 0.6

    def _missile_diversity(self, allocation, targets, missile_types):
        """
        计算导弹组合的多样性得分
        :return: 0(单一类型) - 1(完美混合)
        """
        # 统计各导弹类型的使用比例
        missile_counts = {mt: 0 for mt in missile_types}
        total_used = 0

        for target, mis_allocs in allocation.items():
            for mt, count in mis_allocs.items():
                missile_counts[mt] += count
                total_used += count

        # 计算均衡度 (Entropy-based)
        if total_used == 0:
            return 0.0

        diversity = 0.0
        for count in missile_counts.values():
            p = count / total_used
            if p > 0:
                diversity -= p * math.log(p)

        max_diversity = math.log(len(missile_types))
        return diversity / max_diversity




    def _fallback_allocation(self, targets, target_values, missile_inventory):
        """启发式后备分配方案"""
        allocation = defaultdict(lambda: defaultdict(int))
        remaining_missiles = missile_inventory.copy()

        # 优先分配高价值目标
        for i, target in enumerate(targets):
            if target_values[i] > 0.5:
                for missile_name in remaining_missiles:
                    if remaining_missiles[missile_name] > 0:
                        allocation[target.name][missile_name] = min(2, remaining_missiles[missile_name])
                        remaining_missiles[missile_name] -= allocation[target.name][missile_name]

        # 再分配剩余目标
        for i, target in enumerate(targets):
            if target_values[i] <= 0.5 and remaining_missiles:
                for missile_name in remaining_missiles:
                    if remaining_missiles[missile_name] > 0:
                        allocation[target.name][missile_name] = 1
                        remaining_missiles[missile_name] -= 1
                        break

        return dict(allocation)

    def blue_defense_plan(self, attack_plan: Dict[str, Dict[str, int]]) -> Dict[str, Dict[str, float]]:
        """
        蓝方防御计划
        :param attack_plan: 红方攻击计划 {目标: {导弹类型: 数量}}
        :return: 防御计划 {目标: {导弹类型: 拦截概率}}
        """
        defense_plan = {}

        for target_name, missiles in attack_plan.items():
            # 找到目标舰船
            ship = next(s for s in self.blue_ships if s.name == target_name)

            defense_for_target = {}
            for missile_name, count in missiles.items():
                # 找到对应导弹
                missile = next(m for m in self.red_missiles if m.name == missile_name)

                # 计算此舰船对该导弹的拦截概率
                intercept_prob = ship.intercept_probability(missile)

                defense_for_target[missile_name] = intercept_prob

            defense_plan[target_name] = defense_for_target

        return defense_plan

    def simulate_battle(self, attack_plan: Dict[str, Dict[str, int]], defense_plan: Dict[str, Dict[str, float]]):
        """模拟战斗结果"""
        battle_report = {
            'targets': {},
            'remaining_missiles': {m.name: m.quantity for m in self.red_missiles},
            'ship_status': {s.name: {'damage': 0.0, 'status': 'operational'} for s in self.blue_ships}
        }

        # 模拟每个目标的战斗
        for target_name, missiles in attack_plan.items():
            ship = next(s for s in self.blue_ships if s.name == target_name)

            # 计算该目标预期遭受的总毁伤
            target_damage = 0

            for missile_name, count in missiles.items():
                # 更新剩余导弹
                missile = next(m for m in self.red_missiles if m.name == missile_name)
                battle_report['remaining_missiles'][missile_name] -= count

                # 获取拦截概率
                intercept_prob = defense_plan.get(target_name, {}).get(missile_name, 0.0)

                # 模拟拦截结果
                for _ in range(count):
                    # 判断是否被拦截
                    intercepted = random.random() < intercept_prob

                    if not intercepted:
                        # 未被拦截则造成伤害
                        damage = random.uniform(0.7, 1.3) * missile.warhead_yield
                        target_damage += damage

            # 记录目标损伤
            battle_report['targets'][target_name] = target_damage
            battle_report['ship_status'][target_name]['damage'] += target_damage

            # 判断舰船状态
            if target_damage > 5000:  # 严重损伤阈值(kg TNT)
                battle_report['ship_status'][target_name]['status'] = 'destroyed'
            elif target_damage > 2000:
                battle_report['ship_status'][target_name]['status'] = 'damaged'

        return battle_report


# 示例使用
if __name__ == "__main__":
    # ========== 初始化火力分配系统 ==========
    fas = FireAllocationSystem()

    # ========== 创建红方导弹系统 ==========
    # FJDDDD (ASBM)
    XXX1D = RedMissile(
        name="XXX1D",
        missile_type="ASBM",
        position=(36.0, 123.5),
        rcs=0.1,
        speed=8.0,
        altitude=60000,
        warhead_yield=1000,
        quantity=50
    )

    # CYSXHDD
    XXX2 = RedMissile(
        name="XX-12",
        missile_type="Supersonic",
        position=(36.1, 123.4),
        rcs=0.3,
        speed=3.5,
        altitude=50,
        warhead_yield=250,
        quantity=50
    )

    # YYSDD
    XX02 = RedMissile(
        name="X-XX2",
        missile_type="Subsonic",
        position=(36.2, 123.3),
        rcs=0.6,
        speed=0.9,
        altitude=100,
        warhead_yield=190,
        quantity=50
    )

    fas.add_red_missile(XXX1D)
    fas.add_red_missile(XXX2)
    fas.add_red_missile(XX02)

    # ========== 创建蓝方舰船 ==========
    # 航空母舰
    carrier = BlueShip("CVN-75", "carrier", (35.0, 138.0))
    carrier.set_radar_params(350, 5.4, 500)  # 雷达参数
    carrier.set_aaw_systems([  # 防空系统
        {"type": "CIWS", "count": 1},   #, "range": 3, "reload_time": 5},
        {"type": "SAM", "count": 1},#, "range": 50, "reload_time": 20},
    ])

    # 驱逐舰
    destroyer = BlueShip("DDG-115", "destroyer", (35.1, 138.1))
    destroyer.set_radar_params(280, 6.0, 350)
    destroyer.set_aaw_systems([
        {"type": "CIWS", "count": 1, "range": 3, "reload_time": 5},
        {"type": "SAM", "count": 2, "range": 40, "reload_time": 15},
        {"type": "ECM", "count": 1, "range": 30, "reload_time": 10}
    ])

    # 补给舰
    supply = BlueShip("AOE-6", "supply", (35.2, 138.0))
    supply.set_radar_params(120, 3.0, 100)
    supply.set_aaw_systems([
        {"type": "CIWS", "count": 1, "range": 2, "reload_time": 10}
    ])
    # 巡洋舰（区域防空核心）
    cruiser = BlueShip("CG-67", "cruiser", (35.05, 138.02))
    cruiser.set_radar_params(400, 5.0, 600)  # 高性能相控阵雷达
    cruiser.set_aaw_systems([
        {"type": "SAM", "count": 6, "range": 150, "reload_time": 12},  # 标准-2导弹
        {"type": "SAM", "count": 2, "range": 40, "reload_time": 8},  # 海麻雀导弹
        {"type": "CIWS", "count": 2, "range": 2.5, "reload_time": 4},  # 海拉姆系统
        {"type": "ECM", "count": 1, "range": 45, "reload_time": 6}  # 电子对抗
    ])
    # 护卫舰（近海防御）
    frigate = BlueShip("FFG-62", "frigate", (35.15, 138.05))
    frigate.set_radar_params(180, 6.5, 200)  # 近程雷达
    frigate.set_aaw_systems([
        {"type": "SAM", "count": 1, "range": 30, "reload_time": 20},  # 海麻雀导弹
        {"type": "CIWS", "count": 1, "range": 2, "reload_time": 6},  # 30mm近防炮
        {"type": "ECM", "count": 1, "range": 25, "reload_time": 12}  # 基础电子对抗
    ])
    # 两栖攻击舰（准航母）
    amphibious = BlueShip("LHA-7", "amphibious", (34.98, 138.12))
    amphibious.set_radar_params(350, 4.5, 500)  # 中远程雷达
    amphibious.set_aaw_systems([
        {"type": "SAM", "count": 4, "range": 120, "reload_time": 15},  # 海麻雀导弹
        {"type": "SAM", "count": 2, "range": 25, "reload_time": 10},  # 滚体导弹
        {"type": "CIWS", "count": 2, "range": 3, "reload_time": 5}  # 密集阵
    ])
    # 'CIWS': 0.25,  # 近防炮系统权重
    # 'SAM': 0.60,  # 舰空导弹权重
    # 'ECM': 0.15
    fas.add_blue_ship(carrier)
    fas.add_blue_ship(destroyer)
    fas.add_blue_ship(supply)
    fas.add_blue_ship(cruiser)
    fas.add_blue_ship(frigate)
    fas.add_blue_ship(amphibious)
    ##########################################################
    #fas.set_cost_weight(0.3)  # 较低成本权重，更注重毁伤效果
    #fas.set_cost_weight(0.8)  # 较高成本权重，更注重节省弹药
    # ========== 生成红方火力分配方案 ==========
    print("\n===== 红方火力分配方案 =====")
    start_time = time.time()
    fire_plan = fas.allocate_red_firepower()
    print(f"分配计算时间: {time.time() - start_time:.4f}秒")

    for target, missiles in fire_plan.items():
        print(f"\n目标 {target}:")
        for missile_type, count in missiles.items():
            print(f"  - 分配 {missile_type} 导弹 {count} 枚")
